package org.flexlb.balance.scheduler;

import org.flexlb.balance.strategy.LoadBalancer;
import org.flexlb.config.WhaleMasterConfig;
import org.flexlb.dao.RequestContext;
import org.flexlb.dao.loadbalance.MasterRequest;
import org.flexlb.dao.loadbalance.MasterResponse;
import org.flexlb.dao.loadbalance.ServerStatus;
import org.flexlb.dao.loadbalance.StrategyErrorType;
import org.flexlb.dao.route.RoleType;
import org.flexlb.domain.balance.BalanceContext;
import org.flexlb.enums.LoadBalanceStrategyEnum;
import org.flexlb.service.config.ConfigService;
import org.flexlb.sync.status.EngineWorkerStatus;
import org.flexlb.sync.status.ModelWorkerStatus;
import org.flexlb.trace.WhaleSpan;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.anyString;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.isNull;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class DefaultSchedulerTest {

    @Mock
    private ConfigService configService;

    @Mock
    private WhaleMasterConfig loadBalanceConfig;

    @Mock
    private LoadBalancer prefillLoadBalancer;

    @Mock
    private LoadBalancer decodeLoadBalancer;

    @Mock
    private LoadBalancer vitLoadBalancer;

    @Mock
    private LoadBalancer fusionLoadBalancer;

    @Mock
    private BalanceContext balanceContext;

    @Mock
    private MasterRequest masterRequest;

    @Mock
    private RequestContext requestContext;

    @Mock
    private WhaleSpan span;

    @Mock
    private ModelWorkerStatus modelWorkerStatus;

    private DefaultScheduler defaultScheduler;

    @BeforeEach
    void setUp() {
        // Mock config service
        when(configService.loadBalanceConfig()).thenReturn(loadBalanceConfig);
        when(loadBalanceConfig.getLoadBalanceStrategy()).thenReturn(LoadBalanceStrategyEnum.RANDOM);

        // Create scheduler instance
        defaultScheduler = new DefaultScheduler(configService);

        // Mock LoadBalanceStrategyFactory to return our mock load balancers
        mockStaticLoadBalanceStrategyFactory();

        // Mock balance context
        lenient().when(balanceContext.getMasterRequest()).thenReturn(masterRequest);
        lenient().when(balanceContext.getRequestContext()).thenReturn(requestContext);
        lenient().when(requestContext.getSpan()).thenReturn(span);
        lenient().when(balanceContext.getInterRequestId()).thenReturn(12345L);

        // Mock master request
        lenient().when(masterRequest.getModel()).thenReturn("testModel");

        // Mock model worker status
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PREFILL);
        lenient().when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        // Set up static map for testing
        Map<String, ModelWorkerStatus> testMap = new HashMap<>();
        testMap.put("testModel", modelWorkerStatus);
        EngineWorkerStatus.MODEL_ROLE_WORKER_STATUS_MAP.clear();
        EngineWorkerStatus.MODEL_ROLE_WORKER_STATUS_MAP.putAll(testMap);
    }

    // Helper method to mock the static LoadBalanceStrategyFactory
    private void mockStaticLoadBalanceStrategyFactory() {
        try {
            // Use reflection to set the private fields in DefaultScheduler
            Field prefillField = DefaultScheduler.class.getDeclaredField("prefillLoadBalancer");
            prefillField.setAccessible(true);
            prefillField.set(defaultScheduler, prefillLoadBalancer);

            Field decodeField = DefaultScheduler.class.getDeclaredField("decodeLoadBalancer");
            decodeField.setAccessible(true);
            decodeField.set(defaultScheduler, decodeLoadBalancer);

            Field vitField = DefaultScheduler.class.getDeclaredField("vitLoadBalancer");
            vitField.setAccessible(true);
            vitField.set(defaultScheduler, vitLoadBalancer);

            Field fusionField = DefaultScheduler.class.getDeclaredField("fusionLoadBalancer");
            fusionField.setAccessible(true);
            fusionField.set(defaultScheduler, fusionLoadBalancer);
        } catch (Exception e) {
            fail("Failed to mock LoadBalanceStrategyFactory: " + e.getMessage());
        }
    }

    @Test
    void should_create_scheduler_with_non_null_prefillLoadBalancer_when_initialized() {
        assertNotNull(defaultScheduler, "DefaultScheduler should be created");
        assertNotNull(defaultScheduler.getPrefillLoadBalancer(), "prefillLoadBalancer should not be null");
    }

    @Test
    void should_return_response_with_no_available_worker_error_when_worker_status_map_is_empty() {
        // Setup
        EngineWorkerStatus.MODEL_ROLE_WORKER_STATUS_MAP.clear();

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertNotNull(response, "Response should not be null");
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_AVAILABLE_WORKER.getErrorCode(), response.getCode(), "Error code should match NO_AVAILABLE_WORKER");
        // Note: The method logs an error but doesn't fail when map is empty
    }

    @Test
    void should_return_response_with_no_available_worker_error_when_model_not_in_worker_status_map() {
        // Setup
        when(masterRequest.getModel()).thenReturn("nonExistentModel");

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertNotNull(response, "Response should not be null");
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_AVAILABLE_WORKER.getErrorCode(), response.getCode(), "Error code should match NO_AVAILABLE_WORKER");
        // Note: The method logs an error but doesn't fail when model is missing
    }

    @Test
    void should_return_success_response_with_prefill_and_decode_servers_when_prefill_selection_succeeds() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PREFILL);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus prefillServerStatus = new ServerStatus();
        prefillServerStatus.setSuccess(true);
        prefillServerStatus.setServerIp("192.168.1.1");
        prefillServerStatus.setHttpPort(8080);
        prefillServerStatus.setGroup("group1");
        when(prefillLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PREFILL), isNull())).thenReturn(prefillServerStatus);

        ServerStatus decodeServerStatus = new ServerStatus();
        decodeServerStatus.setSuccess(true);
        decodeServerStatus.setServerIp("192.168.1.2");
        decodeServerStatus.setHttpPort(8081);
        when(decodeLoadBalancer.select(any(BalanceContext.class), eq(RoleType.DECODE), anyString())).thenReturn(decodeServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(2, response.getServerStatus().size(), "Should have 2 server statuses");
        verify(span).setAttribute("prefill_server_ip", "192.168.1.1");
        verify(span).setAttribute("decode_server_ip", "192.168.1.2");
    }

    @Test
    void should_return_response_with_no_prefill_worker_error_when_prefill_selection_fails() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PREFILL);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus prefillServerStatus = new ServerStatus();
        prefillServerStatus.setSuccess(false);
        prefillServerStatus.setMessage("No prefill worker available");
        when(prefillLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PREFILL), isNull())).thenReturn(prefillServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_PREFILL_WORKER.getErrorCode(), response.getCode(), "Error code should match");
        assertNotNull(response.getErrorCode(), "Error message should not be null");
    }

    @Test
    void should_return_success_response_with_fusion_server_when_pdfusion_selection_succeeds() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PDFUSION);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus fusionServerStatus = new ServerStatus();
        fusionServerStatus.setSuccess(true);
        fusionServerStatus.setServerIp("192.168.1.3");
        fusionServerStatus.setHttpPort(8082);
        fusionServerStatus.setGroup("group2");
        fusionServerStatus.setInterRequestId(54321L);
        when(fusionLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PDFUSION), isNull())).thenReturn(fusionServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(1, response.getServerStatus().size(), "Should have 1 server status");
        assertEquals(54321L, response.getInterRequestId(), "Inter request ID should match");
        verify(span).setAttribute("fusion_server_ip", "192.168.1.3");
    }

    @Test
    void should_return_response_with_no_pdfusion_worker_error_when_pdfusion_selection_fails() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PDFUSION);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus fusionServerStatus = new ServerStatus();
        fusionServerStatus.setSuccess(false);
        fusionServerStatus.setMessage("No fusion worker available");
        when(fusionLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PDFUSION), isNull())).thenReturn(fusionServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_PDFUSION_WORKER.getErrorCode(), response.getCode(), "Error code should match");
        assertNotNull(response.getErrorCode(), "Error message should not be null");
    }

    @Test
    void should_return_success_response_with_fusion_and_vit_servers_when_both_selections_succeed() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PDFUSION);
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus fusionServerStatus = new ServerStatus();
        fusionServerStatus.setSuccess(true);
        fusionServerStatus.setServerIp("192.168.1.3");
        fusionServerStatus.setHttpPort(8082);
        fusionServerStatus.setGroup("group2");
        when(fusionLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PDFUSION), isNull())).thenReturn(fusionServerStatus);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(true);
        vitServerStatus.setServerIp("192.168.1.4");
        vitServerStatus.setHttpPort(8083);
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), anyString())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(2, response.getServerStatus().size(), "Should have 2 server statuses");
        verify(span).setAttribute("fusion_server_ip", "192.168.1.3");
        verify(span).setAttribute("vit_server_ip", "192.168.1.4");
    }

    @Test
    void should_return_response_with_no_vit_worker_error_when_vit_selection_fails() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PDFUSION);
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus fusionServerStatus = new ServerStatus();
        fusionServerStatus.setSuccess(true);
        fusionServerStatus.setServerIp("192.168.1.3");
        fusionServerStatus.setHttpPort(8082);
        fusionServerStatus.setGroup("group2");
        when(fusionLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PDFUSION), isNull())).thenReturn(fusionServerStatus);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(false);
        vitServerStatus.setMessage("No vit worker available");
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), anyString())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_VIT_WORKER.getErrorCode(), response.getCode(), "Error code should match");
        assertNotNull(response.getErrorCode(), "Error message should not be null");
    }

    @Test
    void should_log_error_when_master_request_is_null() {
        // Setup
        when(balanceContext.getMasterRequest()).thenReturn(null);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertNotNull(response, "Response should not be null");
    }

    @Test
    void should_return_response_with_no_decode_worker_error_and_release_prefill_cache_when_decode_selection_fails() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PREFILL);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus prefillServerStatus = new ServerStatus();
        prefillServerStatus.setSuccess(true);
        prefillServerStatus.setServerIp("192.168.1.1");
        prefillServerStatus.setHttpPort(8080);
        prefillServerStatus.setGroup("group1");
        when(prefillLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PREFILL), isNull())).thenReturn(prefillServerStatus);

        ServerStatus decodeServerStatus = new ServerStatus();
        decodeServerStatus.setSuccess(false);
        decodeServerStatus.setMessage("No decode worker available");
        when(decodeLoadBalancer.select(any(BalanceContext.class), eq(RoleType.DECODE), anyString())).thenReturn(decodeServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_DECODE_WORKER.getErrorCode(), response.getCode(), "Error code should match NO_DECODE_WORKER");
        verify(prefillLoadBalancer).releaseLocalCache(eq("testModel"), eq("192.168.1.1:8080"), any());
    }

    @Test
    void should_return_success_response_with_vit_server_when_only_vit_role_exists_and_selection_succeeds() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(true);
        vitServerStatus.setServerIp("192.168.1.5");
        vitServerStatus.setHttpPort(8084);
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), isNull())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(1, response.getServerStatus().size(), "Should have 1 server status");
        verify(span).setAttribute("vit_server_ip", "192.168.1.5");
    }

    @Test
    void should_return_response_with_no_vit_worker_error_when_only_vit_role_exists_and_selection_fails() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(false);
        vitServerStatus.setMessage("No vit worker available");
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), isNull())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertFalse(response.isSuccess(), "Response should not be successful");
        assertEquals(StrategyErrorType.NO_VIT_WORKER.getErrorCode(), response.getCode(), "Error code should match");
        assertNotNull(response.getErrorCode(), "Error message should not be null");
    }

    @Test
    void should_return_success_response_with_pdfusion_and_vit_servers_when_both_selections_succeed() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PDFUSION);
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus fusionServerStatus = new ServerStatus();
        fusionServerStatus.setSuccess(true);
        fusionServerStatus.setServerIp("192.168.1.3");
        fusionServerStatus.setHttpPort(8082);
        fusionServerStatus.setGroup("group2");
        fusionServerStatus.setInterRequestId(54321L);
        when(fusionLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PDFUSION), isNull())).thenReturn(fusionServerStatus);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(true);
        vitServerStatus.setServerIp("192.168.1.4");
        vitServerStatus.setHttpPort(8083);
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), anyString())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(2, response.getServerStatus().size(), "Should have 2 server statuses");
        assertEquals(54321L, response.getInterRequestId(), "Inter request ID should match");
        verify(span).setAttribute("fusion_server_ip", "192.168.1.3");
        verify(span).setAttribute("vit_server_ip", "192.168.1.4");
    }

    @Test
    void should_return_success_response_with_prefill_decode_and_vit_servers_when_all_selections_succeed() {
        // Setup
        List<RoleType> roleTypes = new ArrayList<>();
        roleTypes.add(RoleType.PREFILL);
        roleTypes.add(RoleType.VIT);
        when(modelWorkerStatus.getRoleTypeList()).thenReturn(roleTypes);

        ServerStatus prefillServerStatus = new ServerStatus();
        prefillServerStatus.setSuccess(true);
        prefillServerStatus.setServerIp("192.168.1.1");
        prefillServerStatus.setHttpPort(8080);
        prefillServerStatus.setGroup("group1");
        when(prefillLoadBalancer.select(any(BalanceContext.class), eq(RoleType.PREFILL), isNull())).thenReturn(prefillServerStatus);

        ServerStatus decodeServerStatus = new ServerStatus();
        decodeServerStatus.setSuccess(true);
        decodeServerStatus.setServerIp("192.168.1.2");
        decodeServerStatus.setHttpPort(8081);
        when(decodeLoadBalancer.select(any(BalanceContext.class), eq(RoleType.DECODE), anyString())).thenReturn(decodeServerStatus);

        ServerStatus vitServerStatus = new ServerStatus();
        vitServerStatus.setSuccess(true);
        vitServerStatus.setServerIp("192.168.1.5");
        vitServerStatus.setHttpPort(8084);
        when(vitLoadBalancer.select(any(BalanceContext.class), eq(RoleType.VIT), anyString())).thenReturn(vitServerStatus);

        // Execute
        MasterResponse response = defaultScheduler.select(balanceContext);

        // Verify
        assertTrue(response.isSuccess(), "Response should be successful");
        assertNotNull(response.getServerStatus(), "Server status list should not be null");
        assertEquals(3, response.getServerStatus().size(), "Should have 3 server statuses");
        verify(span).setAttribute("prefill_server_ip", "192.168.1.1");
        verify(span).setAttribute("decode_server_ip", "192.168.1.2");
        verify(span).setAttribute("vit_server_ip", "192.168.1.5");
    }
}